[深度学习邪修] Flash Attention 是怎么优化计算过程的🤔

Attention 计算存在的问题

A t t e n t i o n (Q, K, V) = SoftMax (\frac{Q K^{T}}{\sqrt{d}}) V

$QK^T$ $k$ $r$ $Softmax$ $r$ $n$ 个时间片。另外，计算的速度是非常快的，计算的所用时间仅仅占时间片的一小部分，但计算过程中往往伴随着频繁的开销巨大的内存访问，内存访问是高耗时操作，这带来了内存读写瓶颈

得分矩阵计算时的内存操作

$r$ 的所有元素被计算完成，另外，后续计算中同样伴随着多次内存操作:

后续的SoftMax和矩阵乘法操作，每次读取相似度矩阵的一行数据，用于计算Softmax，并参与后续的矩阵乘法

不存储中间结果 $Q、K、V$ 的计算一气呵成。假设我们有某种方法使得Softmax的计算不再阻塞并行过程：

图上所示，假设有一个类似Softmax的计算，但他不依赖一整个向量的数据

为了避免混淆实现与理论，先使用下图的计算过程：

$f$ $f$ 不会阻塞并行计算过程，则最终结果矩阵A的行向量计算公式为:

A_{i} = \sum_{j = 1}^{m} f (q_{i} k_{j}^{T}) v_{j}

$f$ ，该函数不会像Sofxmax那样阻塞并行计算过程。

如标题暗示，Softmax是一个暗含风险的操作，如果一个vector中含有极大值，那么在Sofxmax计算中很容易因为精度问题导致计算结果不准确，并且在后续的量化过程中易导致“饱和”现象发生。Safe-Softmax的改进思想很简单，将极大值从vector中减掉就好了:)

\begin{aligned} S_{s a f e} (x_{i}) & = \frac{e^{x_{i} - m a x (x)}}{\sum e^{x_{j} - m a x (x)}} \\ 等 价 于 \\ S (x_{i}) & = \frac{e^{x_{i}}}{\sum e^{x_{j}}} \end{aligned}

$S_{safe}$ $\sum e^{x_j}$ , 我们总共要进行两次遍历。

有什么方法能够将求最大值操作和求和操作压缩在一起呢？Online-Softmax提出了一种方法：

\begin{aligned} f o r x_{i} & i n V e c t o r [0. . m] : \\ 1. m_{0} & = x_{0} = m a x (- i n f, x_{0}) \\ d_{0} & = e^{x_{0} - m_{0}} \\ 2. m_{1} & = m a x (m_{0}, x_{1}) \\ d_{1} & = e^{x_{0} - m_{1}} + e^{x_{1} - m_{1}} \\ = e^{x_{0} - m_{0} + m_{0} - m_{1}} + e^{x_{1} - m_{1}} \\ = d_{0} \cdot e^{m_{0} - m_{1}} + e^{x_{1} - m_{1}} \\ 3. m_{2} & = m a x (m_{1}, x_{2}) \\ d_{2} & = e^{x_{0} - m_{2}} + e^{x_{1} - m_{2}} + e^{x_{2} - m_{2}} \\ = e^{x_{0} - m_{1} + m_{1} - m_{2}} + e^{x_{1} - m_{1} + m_{1} - m_{2}} + e^{x_{2} - m_{2}} \\ = (e^{x_{0} - m_{1}} + e^{x_{1} - m_{1}}) e^{m_{1} - m_{2}} + e^{x_{2} - m_{2}} \\ = d_{1} \cdot e^{m_{1} - m_{2}} + e^{x_{2} - m_{2}} \\ n. m_{n} & = m a x (m_{n - 1}, x_{n}) \\ d_{n} & = d_{n - 1} \cdot e^{m_{n - 1} - m_{n}} + e^{x_{n} - m_{n}} \end{aligned}

see? 上述计算过程将求和与求最大值操作合并在一个循环当中了！

$d_n$ $m$ 步计算

现在我们只是通过一次遍历完成了最大值和求和操作，最大值用于计算Safe-Softmax，求和作为Safe-Softmax的分母。现在我们既没有计算vector中每个位置的softmax值，也没有计算之后其与value的乘积，此时有下列公式：

\begin{aligned} f o r x_{i} & i n V e c t o r [0. . m], f o r v_{i} i n V a l u e [0. . m] : \\ 1. m_{0} & = x_{0} = m a x (- i n f, x_{0}) \\ d_{0} & = e^{x_{0} - m_{0}} = 1 \\ o_{0} & = \frac{e^{x_{0} - m_{0}}}{d_{0}} \cdot v_{0} \\ 2. m_{1} & = m a x (m_{0}, x_{1}) \\ d_{1} & = d_{0} \cdot e^{m_{0} - m_{1}} + e^{x_{1} - m_{1}} = e^{m_{0} - m_{1}} + e^{x_{1} - m_{1}} \\ o_{1} & = \frac{e^{x_{0} - m_{1}}}{d_{1}} \cdot v_{0} + \frac{e^{x_{1} - m_{1}}}{d_{1}} \cdot v_{1} \\ = \frac{(e^{x_{0} - m_{0} + m_{0} - m_{1}}) d_{0}}{d_{0} \cdot d_{1}} \cdot v_{0} + \frac{e^{x_{1} - m_{1}}}{d_{1}} \cdot v_{1} \\ = o_{0} \cdot \frac{d_{0} e^{m_{0} - m_{1}}}{d_{1}} + \frac{e^{x_{1} - m_{1}}}{d_{1}} \cdot v_{1} \\ s o, o_{j} & = o_{j - 1} \frac{d_{j - 1} e^{m_{j - 1} - m_{j}}}{d_{j}} + \frac{e^{x_{j} - m_{j}}}{d_{j}} \cdot v_{j} \end{aligned}

$j == m$ $A$ 中的单个元素被计算完成。

$r_{i,j}$ $r_i$ $\text{Online-Softmax} * value$ 算法，在单次遍历结束后就能计算出结果矩阵中的一个元素。理论上实现了少量内存访问的并行softmax计算。👍